import os.path

import pandas as pd
from autogluon.common.utils.utils import setup_outputdir
from autogluon.core.utils.loaders import load_pkl
from autogluon.core.utils.savers import save_pkl
from autogluon.tabular import TabularDataset, TabularPredictor


class MultilabelPredictor():
    """
    Tabular Predictor for predicting single or multiple columns in table.
    Creates one or multiple TabularPredictor objects which you can also use individually.
    You can access the TabularPredictor for a particular label via: `multilabel_predictor.get_predictor(label_i)`

    Parameters
    ----------
    labels : List[str]
        The ith element of this list is the column (i.e. `label`) predicted by the ith TabularPredictor stored in this object.
        Can be a single-element list for single-label prediction.
    path : str, default = None
        Path to directory where models and intermediate outputs should be saved.
        If unspecified, a time-stamped folder called "AutogluonModels/ag-[TIMESTAMP]" will be created in the working directory to store all models.
        Note: To call `fit()` twice and save all results of each fit, you must specify different `path` locations or don't specify `path` at all.
        Otherwise files from first `fit()` will be overwritten by second `fit()`.
        Caution: when predicting many labels, this directory may grow large as it needs to store many TabularPredictors.
    problem_types : List[str], default = None
        The ith element is the `problem_type` for the ith TabularPredictor stored in this object.
    eval_metrics : List[str], default = None
        The ith element is the `eval_metric` for the ith TabularPredictor stored in this object.
    consider_labels_correlation : bool, default = True
        Whether the predictions of multiple labels should account for label correlations or predict each label independently of the others.
        If True, the ordering of `labels` may affect resulting accuracy as each label is predicted conditional on the previous labels appearing earlier in this list (i.e. in an auto-regressive fashion).
        Set to False if during inference you may want to individually use just the ith TabularPredictor without predicting all the other labels.
        For single-label prediction, this parameter is ignored.
    kwargs :
        Arguments passed into the initialization of each TabularPredictor.
    """

    multi_predictor_file = 'multilabel_predictor.pkl'

    def __init__(self, labels, path=None, problem_types=None, eval_metrics=None, consider_labels_correlation=True,
                 **kwargs):
        # 修改：允许单标签预测
        if len(labels) < 1:
            raise ValueError("至少需要一个预测标签")

        if (problem_types is not None) and (len(problem_types) != len(labels)):
            raise ValueError("如果提供，`problem_types` 必须与 `labels` 长度相同")
        if (eval_metrics is not None) and (len(eval_metrics) != len(labels)):
            raise ValueError("如果提供，`eval_metrics` 必须与 `labels` 长度相同")

        self.path = setup_outputdir(path, warn_if_exist=False)
        self.labels = labels
        self.is_multilabel = len(labels) > 1  # 新增：标识是否为多标签
        self.consider_labels_correlation = consider_labels_correlation if self.is_multilabel else False

        self.predictors = {}  # key = label, value = TabularPredictor or str path to the TabularPredictor for this label

        if eval_metrics is None:
            self.eval_metrics = {}
        else:
            self.eval_metrics = {labels[i]: eval_metrics[i] for i in range(len(labels))}

        problem_type = None
        eval_metric = None

        for i in range(len(labels)):
            label = labels[i]
            path_i = self.path + "Predictor_" + label

            if problem_types is not None:
                problem_type = problem_types[i]
            if eval_metrics is not None:
                eval_metric = eval_metrics[i]

            # 创建预测器
            self.predictors[label] = TabularPredictor(
                label=label,
                problem_type=problem_type,
                eval_metric=eval_metric,
                path=path_i,
                **kwargs
            )

    def fit(self, train_data, tuning_data=None, **kwargs):
        """
        Fits a separate TabularPredictor to predict each of the labels.

        Parameters
        ----------
        train_data, tuning_data : str or autogluon.tabular.TabularDataset or pd.DataFrame
            See documentation for `TabularPredictor.fit()`.
        kwargs :
            Arguments passed into the `fit()` call for each TabularPredictor.
        """
        if isinstance(train_data, str):
            train_data = TabularDataset(train_data)
        if tuning_data is not None and isinstance(tuning_data, str):
            tuning_data = TabularDataset(tuning_data)

        train_data_og = train_data.copy()
        if tuning_data is not None:
            tuning_data_og = tuning_data.copy()
        else:
            tuning_data_og = None

        save_metrics = len(self.eval_metrics) == 0

        for i in range(len(self.labels)):
            label = self.labels[i]
            predictor = self.get_predictor(label)

            # 修改：处理单标签情况
            if self.is_multilabel:
                if not self.consider_labels_correlation:
                    labels_to_drop = [l for l in self.labels if l != label]
                else:
                    labels_to_drop = [self.labels[j] for j in range(i + 1, len(self.labels))]

                train_data = train_data_og.drop(labels_to_drop, axis=1)
                if tuning_data is not None:
                    tuning_data = tuning_data_og.drop(labels_to_drop, axis=1)
            else:
                # 单标签情况不需要删除任何列
                train_data = train_data_og
                if tuning_data is not None:
                    tuning_data = tuning_data_og

            print(f"Fitting TabularPredictor for label: {label} ...")
            predictor.fit(train_data=train_data, tuning_data=tuning_data, **kwargs)
            self.predictors[label] = predictor.path

            if save_metrics:
                self.eval_metrics[label] = predictor.eval_metric

        self.save()
        return self

    def predict(self, data, **kwargs):
        """
        Returns DataFrame with label columns containing predictions for each label.

        Parameters
        ----------
        data : str or autogluon.tabular.TabularDataset or pd.DataFrame
            Data to make predictions for. If label columns are present in this data, they will be ignored.
            See documentation for `TabularPredictor.predict()`.
        kwargs :
            Arguments passed into the predict() call for each TabularPredictor.
        """
        return self._predict(data, as_proba=False, **kwargs)

    def predict_proba(self, data, **kwargs):
        """
        Returns dict where each key is a label and the corresponding value is the `predict_proba()` output for just that label.

        Parameters
        ----------
        data : str or autogluon.tabular.TabularDataset or pd.DataFrame
            Data to make predictions for.
            See documentation for `TabularPredictor.predict()` and `TabularPredictor.predict_proba()`.
        kwargs :
            Arguments passed into the `predict_proba()` call for each TabularPredictor (also passed into a `predict()` call).
        """
        return self._predict(data, as_proba=True, **kwargs)

    def evaluate(self, data, **kwargs):
        """
        Returns dict where each key is a label and the corresponding value is the `evaluate()` output for just that label.

        Parameters
        ----------
        data : str or autogluon.tabular.TabularDataset or pd.DataFrame
            Data to evaluate predictions of all labels for, must contain all labels as columns.
            See documentation for `TabularPredictor.evaluate()`.
        kwargs :
            Arguments passed into the `evaluate()` call for each TabularPredictor (also passed into the `predict()` call).
        """
        data = self._get_data(data)
        eval_dict = {}

        for label in self.labels:
            print(f"Evaluating TabularPredictor for label: {label} ...")
            predictor = self.get_predictor(label)

            # 修改：处理单标签情况
            if self.is_multilabel and self.consider_labels_correlation:
                # 对于多标签且考虑相关性的情况，需要逐步添加预测结果
                data[label] = predictor.predict(data, **kwargs)
                eval_dict[label] = predictor.evaluate(data, **kwargs)
            else:
                # 单标签或不需要考虑相关性的情况
                eval_dict[label] = predictor.evaluate(data, **kwargs)

        return eval_dict

    def save(self):
        """ Save MultilabelPredictor to disk. """
        for label in self.labels:
            if not isinstance(self.predictors[label], str):
                self.predictors[label] = self.predictors[label].path

        save_pkl.save(path=self.path + self.multi_predictor_file, object=self)
        print(f"MultilabelPredictor saved to disk. Load with: MultilabelPredictor.load('{self.path}')")

    @classmethod
    def load(cls, path):
        """ Load MultilabelPredictor from disk `path` previously specified when creating this MultilabelPredictor. """
        path = os.path.expanduser(path)
        if path[-1] != os.path.sep:
            path = path + os.path.sep
        return load_pkl.load(path=path + cls.multi_predictor_file)

    def get_predictor(self, label):
        """ Returns TabularPredictor which is used to predict this label. """
        predictor = self.predictors[label]
        if isinstance(predictor, str):
            return TabularPredictor.load(path=predictor)
        return predictor

    def _get_data(self, data):
        if isinstance(data, str):
            return TabularDataset(data)
        return data.copy()

    def _predict(self, data, as_proba=False, **kwargs):
        data = self._get_data(data)

        # 修改：处理单标签情况
        if not self.is_multilabel:
            label = self.labels[0]
            print(f"Predicting with TabularPredictor for label: {label} ...")
            predictor = self.get_predictor(label)

            if as_proba:
                return {label: predictor.predict_proba(data, as_multiclass=True, **kwargs)}
            else:
                return data[[label]].assign(**{label: predictor.predict(data, **kwargs)})

        # 多标签处理
        if as_proba:
            predproba_dict = {}

        for label in self.labels:
            print(f"Predicting with TabularPredictor for label: {label} ...")
            predictor = self.get_predictor(label)

            if as_proba:
                predproba_dict[label] = predictor.predict_proba(data, as_multiclass=True, **kwargs)

            # 仅当考虑标签相关性时才更新数据
            if self.consider_labels_correlation:
                data[label] = predictor.predict(data, **kwargs)

        if not as_proba:
            if self.consider_labels_correlation:
                return data[self.labels]
            else:
                # 如果不考虑相关性，单独预测每个标签
                preds = {}
                for label in self.labels:
                    predictor = self.get_predictor(label)
                    preds[label] = predictor.predict(data, **kwargs)
                return pd.DataFrame(preds)
        else:
            return predproba_dict

    def leaderboards(self):
        """ Returns leaderboards for each predictor """
        return [self.get_predictor(label).leaderboard() for label in self.labels]

    # 新增：单标签预测的便捷方法
    def predict_single(self, data, **kwargs):
        """
        当只有一个标签时的便捷预测方法
        返回预测值的Series（而不是DataFrame）
        """
        if self.is_multilabel:
            raise ValueError("此方法仅适用于单标签预测器")
        return self.predict(data, **kwargs)[self.labels[0]]

    # 新增：单标签概率预测的便捷方法
    def predict_proba_single(self, data, **kwargs):
        """
        当只有一个标签时的便捷概率预测方法
        返回概率数组（而不是字典）
        """
        if self.is_multilabel:
            raise ValueError("此方法仅适用于单标签预测器")
        return self.predict_proba(data, **kwargs)[self.labels[0]]